import os
import torch
from torch import nn
import argparse

from dataclasses import replace
from http.server import HTTPServer, BaseHTTPRequestHandler
import json,os,re,posixpath,codecs
from urllib.parse import unquote
import cgi
from cgi import parse_header, parse_multipart
from urllib.parse import parse_qs
import numpy as np
import time
from global_config import *
from BSGenerator import BSGenerator
import base64
from txt2json import *

import google.protobuf as pb
print(f'protobuf version: {pb.__version__}')
from YAudio2Bs_pb2 import Audio2BSData
import threading


from concurrent import futures
from google.cloud import pubsub_v1
from typing import Callable

lock = threading.Lock()

generator = BSGenerator(ckp=ckp)   
      
use_google_publisher = False
publisher = None
topic_path = None
       
def jsonToProto(jdata,fileId):
    pdata = Audio2BSData()
    for frame in jdata.frames:
        pframe = pdata.frames.add()
        for location in frame.location:
            pframe.location.append(location)
        for weight in frame.weight:
            pframe.weight.append(weight)
        pframe.time = frame.time
    if fileId is not None:
        pdata.fileId = fileId
    s = pdata.SerializeToString()
    return s 
        
requestCount = 0
class AIHTTPRequestHandler(BaseHTTPRequestHandler):
    def handle_timeout(self):
        self.send_response(408)
        self.send_header('Content-type','text/html')
        self.end_headers()
        self.wfile.write(b'Timeout')
    def do_GET(self):
        data={}
        data['ret'] = "ok"
        jstr = json.dumps(data)   
        self.send_response(500)
        self.send_header('Content-Type', 'application/json')
        self.end_headers()
        self.wfile.write(jstr.encode())
        return
        
    def do_POST(self):
        lock.acquire()
        if self.path == '/':
            self.do_common()
        elif self.path == "/gcp":
            self.do_google()
        
        lock.release()
    def do_google(self):
        self.do_common()
        
    def do_common(self):
        
        
        print('start do post')       
        
        global requestCount
        requestCount += 1       
        
        try:
            postvars = ''
            if use_google_publisher:
                
                print('start get headers content-length')     
                length = int(self.headers['content-length'])    
                print(f'end get headers content-length: {length}')  
                rawdata = self.rfile.read(length)   
                
                strdata = rawdata.decode('utf-8')
                # print(f'end read data {strdata}')   
                # postvars = base64.b64decode(rawdata).decode('utf-8')
                # print(f'end endcode read data {postvars}')  
                postvars = strdata
                
            else:
                print('start parse_POST')  
                postvars = self.parse_POST()        

            print('end parse post param')
            json_object = json.loads(postvars)
            
            base64_str = json_object['WavData']
            
            fileId = None
            if 'fileId' in json_object:
                fileId = json_object['fileId']
            
            if 'UseProtobuf' not in json_object:
                useproto = False
            else:
                useproto = json_object['UseProtobuf']
            
            byte_data = base64.b64decode(base64_str)#将base64转换为二进制
            tempPath = './temp.wav'
            
            file = open(tempPath, 'wb')
            try:
                file.write(byte_data)
            finally:
                file.close()
        
            data = generator.computeBS(wav_file=tempPath)   
            
            if useproto:
                if use_google_publisher:                   
                
                    pbytes = jsonToProto(data,fileId)
                    publish(pbytes)
                    
                    dictdata = dict()
                    dictdata['fileId'] = fileId
                    dictdata['ret'] = 'OK'
                    jstr = json.dumps(dictdata)   
                    self.send_response(200)
                    self.send_header('Content-Type', 'application/json')
                    self.end_headers()
                    data = jstr.encode()
                    self.wfile.write(data)
                else:
                    pbytes = jsonToProto(data,fileId)
                    self.send_response(200)
                    self.send_header('Content-Type', 'application/octet-stream')
                    self.end_headers()
                    print(f"protobuf send data len: {len(pbytes)}")
                    self.wfile.write(pbytes)
                    
                
            else:
                dictdata = asdict(data)
                dictdata['fileId'] = fileId
                dictdata['ret'] = 'OK'
                jstr = json.dumps(dictdata)   
                self.send_response(200)
                self.send_header('Content-Type', 'application/json')
                self.end_headers()
                data = jstr.encode()
                print(f"json send data len: {len(data)}")
                self.wfile.write(data)
                
                publish(data)
            
            print(f'reqcount {requestCount} Post Sucess')
        except Exception as error:
            print(f'reqcount {requestCount} Post Connection Error : {error}')
            
            self.send_response(500)
            self.send_header('Content-Type', 'application/json')
            self.end_headers()
            jstr = '{}'
            data = jstr.encode()
            self.wfile.write(data)
            
        
    

    def do_HEAD(self):
        self._set_headers()

    def parse_POST(self):
        ctype, pdict = parse_header(self.headers['content-type'])
        print(self.headers)

        if ctype == 'multipart/form-data':
            postvars = cgi.parse_multipart(self.rfile, pdict)
        if ctype == 'application/json':
            length = int(self.headers['content-length'])          
            postvars = self.rfile.read(length).decode('utf8')
        elif ctype != 'application/x-www-form-urlencoded':
            length = int(self.headers['content-length'])
            postvars = parse_qs(self.rfile.read(length), keep_blank_values=1)
        else:
            postvars = {}        
        return postvars

def publish(data):    
    if use_google_publisher:
        # When you publish a message, the client returns a future.
        future = publisher.publish(topic_path, data)
        print(f'publish to google topic_path: {topic_path}')        
        print('publish future result:' + future.result())
  
if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument("--use_google", action="store_true")
    args = parser.parse_args()
    use_google_publisher = args.use_google

    if use_google_publisher:
        os.environ['GOOGLE_APPLICATION_CREDENTIALS']= os.path.join(os.path.dirname(os.path.abspath(__file__)), 'yuji-20240819-828675ef42a3.json')
        project_id = "yuji-20240819"
        # topic_id = "projects/yuji-20240819/topics/mouth-data-topic"
        topic_id = "mouth-data-topic"
        publisher = pubsub_v1.PublisherClient()
        # The `topic_path` method creates a fully qualified identifier
        # in the form `projects/{project_id}/topics/{topic_id}`
        topic_path = publisher.topic_path(project_id, topic_id)

    while True:
        try:    
            # httpd = HTTPServer(('localhost', 10001), AIHTTPRequestHandler)
            httpd = HTTPServer(('0.0.0.0', 10001), AIHTTPRequestHandler)
            httpd.timeout = 10
            httpd.serve_forever()       
        except Exception as error:
            print(f'httpserver Error : {error}')
            httpd.socket.close()